package vip.bblog.dynamic.datasource;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;


/**
 * 动态数据源注册
 *
 * @author An-Il
 * @version v.0.1
 */
public class DynamicDataSourceRegister implements InitializingBean, DisposableBean {

    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSourceRegister.class);

    private Map<Object, Object> targetDataSources = new HashMap<>();

    private Object defaultTargetDataSource;

    private DynamicDataSourceProperty dataSourceProperty;

    public DynamicDataSourceRegister() {
    }

    public DynamicDataSourceRegister(DynamicDataSourceProperty dataSourceProperty) {
        this.dataSourceProperty = dataSourceProperty;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        this.registerDataSource();
    }

    /**
     * 注册多数据源
     */
    public void registerDataSource() {
        dataSourceProperty.getDatasource().forEach((key, value) -> {
            if (StringUtils.isEmpty(value.getPoolName())) {
                value.setPoolName(key);
            }
            DynamicDataSourceContextHolder.dataSourceIds.add(key);
            DataSource dataSource = createHikariDataSource(value);
            targetDataSources.put(key, dataSource);
        });
        logger.info("共注册 {} 个数据源", targetDataSources.size());
    }

    /**
     * 创建Hikari的数据源
     *
     * @param hikariConfig
     * @return
     */
    public DataSource createHikariDataSource(HikariConfig hikariConfig) {
        return new HikariDataSource(hikariConfig);
    }


    /**
     * 关闭数据源
     *
     * @throws Exception
     */
    @Override
    public void destroy() {
        logger.info("starting closing dynamicDataSource...");
        targetDataSources.forEach((key, dataSource) -> {
            Class<? extends DataSource> clazz = ((DataSource) dataSource).getClass();
            try {
                Method closeMethod = clazz.getDeclaredMethod("close");
                closeMethod.invoke(dataSource);
            } catch (Exception e) {
                logger.warn("关闭数据源 {} 失败,", key);
            }
        });
    }

    public DynamicDataSourceProperty getDataSourceProperty() {
        return dataSourceProperty;
    }

    public void setDataSourceProperty(DynamicDataSourceProperty dataSourceProperty) {
        this.dataSourceProperty = dataSourceProperty;
    }

    public Map<Object, Object> getTargetDataSources() {
        return targetDataSources;
    }

    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources = targetDataSources;
    }

    public Object getDefaultTargetDataSource() {
        return targetDataSources.get(dataSourceProperty.getPrimary());
    }

    public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
        this.defaultTargetDataSource = defaultTargetDataSource;
    }
}
